package com.dek.grpc.stream;

import io.grpc.Server;
import io.grpc.ServerBuilder;
import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;

import java.io.IOException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Logger;

public class StreamServer {

    private static final Logger logger =
            Logger.getLogger(StreamServer.class.getName());

    public static void main(String[] args) throws InterruptedException, IOException {
        StreamingHelloGrpc.StreamingHelloImplBase streamingHelloImplBase = new StreamingHelloGrpc.StreamingHelloImplBase() {
            @Override
            public StreamObserver<HelloRequest> sayHelloStreaming(StreamObserver<HelloReply> responseObserver) {
                logger.info("client call ....");
                final ServerCallStreamObserver<HelloReply> serverCallStreamObserver =
                        (ServerCallStreamObserver<HelloReply>) responseObserver;
                serverCallStreamObserver.disableAutoInboundFlowControl();

                final AtomicBoolean wasReady = new AtomicBoolean(false);
                serverCallStreamObserver.setOnCancelHandler(() -> {
                    if (serverCallStreamObserver.isReady() && wasReady.compareAndSet(false, true)) {
                        serverCallStreamObserver.request(1);
                    }
                });
                return new StreamObserver<HelloRequest>() {
                    @Override
                    public void onNext(HelloRequest request) {
                        // Accept and enqueue the request.
                        System.out.println(111);
                        try {
                            String name = request.getName();
                            logger.info("--> " + name);

                            // Simulate server "work"
                            Thread.sleep(100);

                            // Send a response.
                            String message = "Hello " + name;
                            logger.info("<-- " + message);
                            HelloReply reply = HelloReply.newBuilder().setMessage(message).build();
                            responseObserver.onNext(reply);

                            // Check the provided ServerCallStreamObserver to see if it is still ready to accept more messages.
                            if (serverCallStreamObserver.isReady()) {
                                // Signal the sender to send another request. As long as isReady() stays true, the server will keep
                                // cycling through the loop of onNext() -> request()...onNext() -> request()... until either the client
                                // runs out of messages and ends the loop or the server runs out of receive buffer space.
                                //
                                // If the server runs out of buffer space, isReady() will turn false. When the receive buffer has
                                // sufficiently drained, isReady() will turn true, and the serverCallStreamObserver's onReadyHandler
                                // will be called to restart the message pump.
                                serverCallStreamObserver.request(1);
                            } else {
                                // If not, note that back-pressure has begun.
                                wasReady.set(false);
                            }
                        } catch (Throwable throwable) {
                            throwable.printStackTrace();
                            responseObserver.onError(
                                    Status.UNKNOWN.withDescription("Error handling request").withCause(throwable).asException());
                        }
                    }

                    @Override
                    public void onError(Throwable t) {
                        // End the response stream if the client presents an error.
                        t.printStackTrace();
                        responseObserver.onCompleted();
                    }

                    @Override
                    public void onCompleted() {
                        // Signal the end of work when the client ends the request stream.
                        logger.info("COMPLETED");
                        responseObserver.onCompleted();
                    }
                };
            }
        };

        final Server server = ServerBuilder
                .forPort(50051)
                .addService(streamingHelloImplBase)
                .build()
                .start();

        logger.info("Listening on " + server.getPort());

        Runtime.getRuntime().addShutdownHook(new Thread() {
            @Override
            public void run() {
                logger.info("Shutting down");
                server.shutdown();
            }
        });
        server.awaitTermination();
    }


}
